/*
 * File:   soah3d.c
 * Author: boyan
 *
 * Created on November 28, 2007, 11:53 AM
 */

#include "vecs.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "glodef.h"
#include "errlog.h"
#include "glovar.h"
#include "grid3d.h"

#include "fe.h"
#include "fe_matr.h"

#include "transfer.h"

#include "f3io.h"

#include "soah_common.h"

#include <petsc.h>


#if defined(USE_QGRID)
#include "g3d_mesq.h"
#endif


/* this is in built.c */
extern const char * const built_str;

/* the following are defined in fem3in.c */
extern const char nrms_fname[];
extern char log_fname[];
extern int dump_bin;
extern void parse_cmd_line(int argc, char *argv[]);

/* the following are defined at the end of this file */
static void set_advection_flags(void);
/*
 * static void set_order_flags(void);
 */

static vector init_u;
static scalar init_p;
static vector finl_v;
static scalar finl_q;
void solve_adjoint_11_fvp(const vector v_finl, scalar q_finl, vector vtT_bc);

int mpi_num_procs, mpi_my_id;

void clenup_soah3d(void)
{
    PrintLog("Cleaning up: \n");
    free_vec(init_u);
    free_vec(finl_v);
    free_scal(init_p);
    free_scal(finl_q);

#if defined(USE_QGRID)
    mesq_done(mesquite);
#endif

    free_glovar();

    ReleaseGrid(&grid);
}

void done_petsc(void)
{
    PetscFinalize();
}

int main(int argc, char** argv)
{

    PetscInitialize(&argc, &argv, 0, 0);
    atexit(done_petsc);

    /* Initialize MPI */

    MPI_Comm_size(PETSC_COMM_WORLD, &mpi_num_procs);
    MPI_Comm_rank(PETSC_COMM_WORLD, &mpi_my_id);

    /*
        if(mpi_num_procs > 1)
        {
            if(mpi_my_id == 0)
            {
                char *fname = strrchr(argv[0], '/');
                fname = fname ? fname+1 : argv[0];
                printf("%s is only available for single processor.\n", fname);
            }
            return EXIT_FAILURE;
        }
     */

    /* Start logging before all else */
    InitLogfile(NULL);
    PrintLog("%s", built_str);

    set_defaults();

    /* parse the command line and input file */
    if(argc>1)
        parse_cmd_line(argc, argv);

    atexit(clenup_soah3d);

    PRINTF("%s", built_str), fflush(stdout);

    /* we cannot live without a grid */
    if(grid==NULL)
        error_msg("no grid\n");
    grid->MAXH=GridSize(grid, grid->nodes);
    grid->MINQ=GridQuality(grid, grid->nodes);
    PrintLog("(INFO) Reference grid: size %g     quality %g\n",
        (double)grid->MAXH, (double)grid->MINQ);
#if defined USE_QGRID
    mesquite=mesq_init(grid);
    PrintLog("(INFO) Running Mesquite on reference grid\n");
    mesq_run_no_wrapper(mesquite);
    grid->MAXH=GridSize(grid, grid->nodes);
    grid->MINQ=GridQuality(grid, grid->nodes);
    PrintLog("(INFO) Reference grid: size %g     quality %g\n",
        (double)grid->MAXH, (double)grid->MINQ);
#endif


    /* we also need a "time frame" */
    if(total_ti<=curr_ti)
        error_msg("no time params\n");

    /* set the treatment of advection terms */
    set_advection_flags();

    /* set the time integration order */
    /* !!! NOTE: do this after the treatment of advection is set !!! */
    /*
     * set_order_flags();
     */

    if(flags.shape_opt)
    {
        flags.need_adjoint=TRUE;
        // if(alegrid==NULL)
        // alegrid=init_alegrid(grid);
    }
    if(flags.periodic==per_armijo||flags.periodic==per_adjoint||
        flags.periodic==per_lbfgs)
    {
        flags.need_adjoint=TRUE;
    }
    if(flags.conv_adjoint) flags.need_adjoint=TRUE;
    if(flags.need_adjoint)
    {
        if(flags.need_adjoint&& !flags.conv_adjoint)
            dump_bin=1;
        if(flags.conv_adjoint||flags.convergence||flags.stationary|| !flags.shape_opt)
            flags.adj_advection=flags.advection;
        else
            flags.adj_advection=flags.advection;
        /*
         * flags.adj_advection = adv_explicit;
         */
        PrintLog("flags.adj_advection = %d\n", flags.adj_advection);
    }


    if(dump_bin>0)
        create_dir(root_dir());

    /* Rename the log file */
#if (1)
    {
        char *s;
        if(log_fname[0]=='\0')
            sprintf(log_fname, "%s.log", grid->gridname);
        s=strrchr(log_fname, '.');
        if(flags.convergence)
            sprintf(s, "_%.20s_dt%g_o%d.log", grid->gridname, (double)TStep,
            flags.time_order);
        else if(flags.conv_adjoint)
            sprintf(s, "_adj_%.20s_dt%g_o%d.log", grid->gridname, (double)TStep,
            flags.time_order);
        else
            sprintf(s, "_dt%g_o%d.log", (double)TStep, flags.time_order);
        if(dump_bin>0)
        {
            create_dir(logs_dir());
            if(!flags.convergence&& !flags.conv_adjoint)
            {
                s=malloc(FILENAME_MAX);
                strcpy(s, log_fname);
                sprintf(log_fname, "%s/%s", logs_dir(), s);
                free(s);
            }
            open_dir(root_dir());
        }
        RenameLogfile(log_fname);
    }
#endif

    init_state_solver(); /* also initializes the adjoint solver if flags.need_adjoint is up */
    alloc_glovar(elemU, elemP);

    if(elemU==fe_bub_vec)
    {
        init_u=alloc_vec(fe_ndof(grid, fe_bub));
        finl_v=alloc_vec(fe_ndof(grid, fe_bub));
    }
    else
    {
        init_u=alloc_vec(fe_ndof(grid, elemU));
        finl_v=alloc_vec(fe_ndof(grid, elemU));
    }
    init_p=alloc_scal(fe_ndof(grid, elemP));
    finl_q=alloc_scal(fe_ndof(grid, elemP));

    /***********************************************************
     *  this is a test until down there 
     */

    // convergence test for sigma integral with state and adjoint
    //      velocity given explicitly
    /*
    int run = 1;
    while(run < 3)
    {
        extern double norm_sigma_vecscal_z(grid3d_t *grid, const vector v1,
                                           const scalar s2);
        extern double norm_sigma_vector(grid3d_t *grid, const vector v1,
                                        const vector v2);

        if( run == 1 )
        {
            create_dir(state_dir());
            create_dir(adjoint_dir());
            for(curr_ti = 0; curr_ti <= total_ti; curr_ti++)
            {
                interpolate_vector_lagr(grid, elemU, TIME, exct_vel, u);
                interpolate_scalar_lagr(grid, elemP, TIME, exct_prs, p);
                make_scal_mean_0(grid, elemP, p);
                save_state_time(state_dir());

                interpolate_vector_lagr(grid, elemU, TIME, adj_exct_vel, v);
                interpolate_scalar_lagr(grid, elemP, TIME, adj_exct_prs, q);
                make_scal_mean_0(grid, elemP, q);
                save_adjoint_time(adjoint_dir());
            }
        }
        else
        {
            // solve the state equation ...
            curr_ti = 0;
            interpolate_vector_lagr(grid, elemU, TIME, init_vel, init_u);
            interpolate_scalar_lagr(grid, elemP, TIME, init_prs, init_p);
            make_scal_mean_0(grid, elemP, init_p);

            solve_state_ivp(init_u, init_p);
            print_conv_norms();

            // solve the adjoint equation ...
            curr_ti = total_ti;
            interpolate_vector_lagr(grid, elemU, TIME, adj_finl_vel, finl_v);
            interpolate_scalar_lagr(grid, elemP, TIME, adj_finl_prs, finl_q);
            make_scal_mean_0(grid, elemP, finl_q);

            extern void solve_adjoint_analytic_load_state_fvp
                    (const vector v_finl, scalar q_finl);

            solve_adjoint_analytic_load_state_fvp(finl_v, finl_q);
            print_adj_conv_norms();
        }

        alloc_sg();

        // now compute the approximate shape gradient
        compute_shape_gradient();

        // BUG FIX: invert the sigma-mass matrix on xi!!!
        extern fe_matr_p sg_mass;
        extern scalar sg_w1, sg_w2;
        char fname[200];

        fe_matr_assemble(grid, 1, &sg_mass);
        
        // sprintf(fname, "%s/sg_mass.txt", root_dir());
        // char fname_bc[200];
        // sprintf(fname_bc, "%s/sg_mass_bc.txt", root_dir());
        // fe_matr_print(sg_mass, fname, fname_bc);
        
        copy_scal(sg_xi, sg_w1);
        fe_matr_solve_scalar(sg_mass, sg_xi, sg_w1, sg_w2, 0.0);

        sprintf(fname, "%s/nodes.txt", root_dir());
        print_nodes(grid, grid->nodes, fname);
        sprintf(fname, "%s/elems.txt", root_dir());
        print_elems(grid, fname);
        sprintf(fname, "%s/xi_fixed_%d.txt", root_dir(), run);
        print_scal(sg_xi, fname);
        // BUG FIX: done

        // compute the exact shape gradient
        vector wv1 = alloc_vec(grid->NN);
        vector wv2 = alloc_vec(grid->NN);
        vector esg = alloc_vec(grid->NN);

        curr_ti = 0;
        apply_bc_dir_vector(grid, fe_p1, TIME, shgrd_vel, wv1);
        for(curr_ti = 1; curr_ti <= total_ti; curr_ti++)
        {
            copy_vec(wv1, wv2);
            apply_bc_dir_vector(grid, fe_p1, TIME, shgrd_vel, wv1);

            axpy_vec(esg, 0.5 * TStep, wv1);
            axpy_vec(esg, 0.5 * TStep, wv2);
        }

        sprintf(fname, "%s/esg.txt", root_dir());
        print_vec(esg, fname);

        double norm = norm_sigma_vecscal_z(grid, esg, sg_xi);
        double norm1 = norm_sigma_vector(grid, esg, NULL);
        double relerr = norm / norm1;
        PRINTF("\r|| esg - xi*z || = %10.8g%30s\n", norm, "");
        PRINTF("\r||     esg    || = %10.8g%30s\n", norm1, "");
        PRINTF("\r  relative error = %5.1f%-30s\n", relerr * 100.0, "%");

        free_vec(esg);
        free_vec(wv1);
        free_vec(wv2);


        free_sg();

        // write values to file
        if( run == 1 )
            sprintf(fname, "sg_exct_uv_norms.dat");
        else
            sprintf(fname, "sg_aprx_uv_norms.dat");
        FILE *F = fopen(fname, "a+t");
        if( F == NULL )
            error_msg("Cannot open file %s for writing.\n", fname);
        fseek(F, 0, SEEK_SET);
        if( fgetc(F) == EOF )
        {
            fprintf(F, "%-12s    %-12s    %-14s  %-14s  %-14s\n",
                    "dT", "grid", "||xi*z-esg||", "||esg||", "relative(%)");
        }
        else
            fseek(F, 0, SEEK_END);
        fprintf(F, "%-12g    %-12g", (double) TStep, (double) grid->MAXH);
        fprintf(F, "    %12e", norm);
        fprintf(F, "    %12e", norm1);
        fprintf(F, "    %12.8g", relerr * 100.0);
        fprintf(F, "\n");
        fclose(F);

        if( ++run == 3 )
            goto cleanup;
    }
     */



    /*
    {
        extern double soah_shape_cost(vector init_u, scalar init_p);
        extern void test_vector_sigma(const grid3d_t *grid, vector u, scalar v);
        extern double norm_sigma_vector(grid3d_t *grid, const vector v1, const vector v2);
        extern double norm_sigma_vecscal_z(grid3d_t *grid, const vector v1, const scalar s2);
        double dt = TStep;
        alloc_sg();
        int k = 1;
        // for(k=0; k<5; k++)
        {

            SetTimeParams(0.0, StopTime, dt / (double) (1 << k));

            // solve the state equation ...
            curr_ti = 0;
            interpolate_vector_lagr(grid, elemU, TIME, init_vel, init_u);
            interpolate_scalar_lagr(grid, elemP, TIME, init_prs, init_p);
            make_scal_mean_0(grid, elemP, init_p);

            soah_shape_cost(init_u, init_p);
            PrintLog("\r k=%d dt=%10.8g F=%10.8g\n", k, TStep, total_cost);
            PRINTF("\r k=%d dt=%10.8g F=%10.8g%25s\n", k, TStep, total_cost, " ");

            // solve the adjoint equation ...
            curr_ti = total_ti;
            interpolate_vector_lagr(grid, elemU, TIME, adj_finl_vel, finl_v);
            interpolate_scalar_lagr(grid, elemP, TIME, adj_finl_prs, finl_q);
            make_scal_mean_0(grid, elemP, finl_q);

            solve_adjoint_analytic_fvp(finl_v, finl_q, NULL);

            char fname[FILENAME_MAX];
            sprintf(fname, "%s/nodes.txt", root_dir());
            print_nodes(grid, grid->nodes, fname);
            sprintf(fname, "%s/elems.txt", root_dir());
            print_elems(grid, fname);

            // now check the shape gradient
            compute_shape_gradient();

            // compute the exact shape gradient
            vector wv1 = alloc_vec(grid->NN);
            vector wv2 = alloc_vec(grid->NN);
            vector esg = alloc_vec(grid->NN);

            curr_ti = 0;
            apply_bc_dir_vector(grid, fe_p1, TIME, shgrd_vel, wv1);
            for(curr_ti = 1; curr_ti <= total_ti; curr_ti++)
            {
                copy_vec(wv1, wv2);
                apply_bc_dir_vector(grid, fe_p1, TIME, shgrd_vel, wv1);

                axpy_vec(esg, 0.5 * TStep, wv1);
                axpy_vec(esg, 0.5 * TStep, wv2);
            }

            sprintf(fname, "%s/esg.txt", root_dir());
            print_vec(esg, fname);


            extern fe_matr_p sg_mass;
            extern scalar sg_w1, sg_w2;
            extern vector sg_z;

            fe_matr_assemble(grid, 1, &sg_mass);

            sprintf(fname, "%s/sg_mass.txt", root_dir());
            char fname_bc[200];
            sprintf(fname_bc, "%s/sg_mass_bc.txt", root_dir());
            fe_matr_print(sg_mass, fname, fname_bc);

            fe_matr_solve_scalar(sg_mass, sg_w1, sg_xi, sg_w2, 0.0);

            sprintf(fname, "%s/xi1.txt", root_dir());
            print_scal(sg_w1, fname);

            double norm = norm_sigma_vecscal_z(grid, esg, sg_w1);
            PRINTF("\r|| esg - xi*z || = %10.8g%50s\n", norm, "");
            double norm1 = norm_sigma_vector(grid, esg, NULL);
            PRINTF("\r||     esg    || = %10.8g%50s\n", norm1, "");
            PRINTF("\r  relative error = %5.1f%-50s\n", norm / norm1 * 100.0, "%");

            test_vector_sigma(grid, esg, sg_xi);
            sprintf(fname, "%s/esg_tested.txt", root_dir());
            print_scal(sg_xi, fname);

            fe_matr_solve_scalar(sg_mass, sg_w1, sg_xi, sg_w2, 0.0);
            sprintf(fname, "%s/xi_solved.txt", root_dir());
            print_scal(sg_w1, fname);

            use_scalar(sg_w1)
            use_vector(sg_z)
            use_vector(wv1)
            for(int i = 0; i < grid->NN; i++)
            {
                wv1X[i] = sg_w1F[i] * sg_zX[i];
                wv1Y[i] = sg_w1F[i] * sg_zY[i];
                wv1Z[i] = sg_w1F[i] * sg_zZ[i];
            }
            done_vector(wv1)
            done_vector(sg_z)
            done_scalar(sg_w1)

            sprintf(fname, "%s/esg_solved.txt", root_dir());
            print_vec(wv1, fname);

            free_vec(esg);
            free_vec(wv1);
            free_vec(wv2);

        }

        free_sg();

        goto cleanup;
    }
     */
    /*
     *  the test ends here
     ***********************************************************/

    if(flags.conv_adjoint||flags.convergence||flags.stationary|| !flags.shape_opt)
    {
        /************************************************************************/
        if(dump_bin>0)
            save_ref_grid();

        if(dump_bin>0&&flags.moving_grid)
        {
            char mesh_log[FILENAME_MAX];
            sprintf(mesh_log, "%s/mesh.log", logs_dir());
            PrintLog("Logging mesh generation to file %s\n", mesh_log);
            GetLogfilePath(log_fname, FILENAME_MAX-1);
            InitLogfile(mesh_log);
            soah_generate_meshes(TRUE);
            AppendLogfile(log_fname);
        }

        if(!flags.conv_adjoint)
        {
            /* set initial condition */
            curr_ti=0;
            move_grid();
            /* TurnLogOff(); */
            interpolate_vector_lagr(grid, elemU, TIME, init_vel, init_u);
#ifndef CS_INITIAL_PRESSURE
            interpolate_scalar_lagr(grid, elemP, TIME, init_prs, init_p);
            make_scal_mean_0(grid, elemP, init_p);
#endif
            /* TurnLogOn(); */
            soah_shape_cost(init_u, init_p);

            PRINTF("\r%60s\r", " ");
            PRINTF("||u-eu||_L2linf = %10.4e\n", time_inf_solution_norm[NORM_ERR_U_L2]);
            PRINTF("||p-ep||_L2linf = %10.4e\n", time_inf_solution_norm[NORM_ERR_P_L2]);
            PRINTF("||u-eu||_H1linf = %10.4e\n", time_inf_solution_norm[NORM_ERR_U_H1]);
            PRINTF("||p-ep||_H1linf = %10.4e\n", time_inf_solution_norm[NORM_ERR_P_H1]);

            if(flags.convergence)
                print_conv_norms();
        }
        else
        {
            /* set final condition */
            curr_ti=total_ti;
            move_grid();
            /* TurnLogOff(); */
            interpolate_vector_lagr(grid, elemU, TIME, adj_finl_vel, finl_v);
            interpolate_scalar_lagr(grid, elemP, TIME, adj_finl_prs, finl_q);
            make_scal_mean_0(grid, elemP, finl_q);
            /* TurnLogOn(); */

            if(!solve_adjoint_analytic_fvp(finl_v, finl_q, NULL));
            PRINTF("\r%60s\r", " ");
            PRINTF("||v-ev||_L2linf = %10.4e\n", time_inf_adj_solution_norm[NORM_ERR_U_L2]);
            PRINTF("||q-eq||_L2linf = %10.4e\n", time_inf_adj_solution_norm[NORM_ERR_P_L2]);
            PRINTF("||v-ev||_H1linf = %10.4e\n", time_inf_adj_solution_norm[NORM_ERR_U_H1]);
            PRINTF("||q-eq||_H1linf = %10.4e\n", time_inf_adj_solution_norm[NORM_ERR_P_H1]);

            if(flags.conv_adjoint)
                print_adj_conv_norms();

        }

        goto cleanup;
        /************************************************************************/
    }
    else /* flags.shape_opt */
    {
        extern void soah_shape_bloat(vector init_u, scalar init_p,
            vector finl_v, scalar finl_q);
        /************************************************************************/

        PrintLog("cost functional with alpha = %5g\tbeta = %5g\n", alpha, beta);
        PRINTF("alpha = %5g\tbeta = %5g\n", alpha, beta);

        clean_vec(init_u);
        clean_vec(finl_v);
        clean_scal(init_p);
        clean_scal(finl_q);
        /*
         soah_shape_optimize(init_u, init_p, finl_v, finl_q);
         */
        // if(alegrid==NULL) alegrid=init_alegrid(grid);

        switch(flags.shape_opt)
        {
            case opt_pit:
                soah_shape_optimize_pit(init_u, init_p, finl_v, finl_q);
                break;
            case opt_armijo:
                soah_shape_optimize_armijo(init_u, init_p, finl_v, finl_q);
                break;
            case opt_custom:
                // soah_shape_bloat(init_u, init_p, finl_v, finl_q);
                soah_shape_optimize_test_gradient_at_point(init_u, init_p, finl_v, finl_q);
                break;
            default:
                error_msg("Unknown shape_opt flag %d\n", flags.shape_opt);

        }

        goto cleanup;
        /************************************************************************/
    }

cleanup:
    {
        PrintLogLine(78, '=');

        CloseLogfile();
        PRINTF("\n");
    }
    return(EXIT_SUCCESS);
}

/****************************************************************************/

static void set_advection_flags(void)
{
    if(flags.advection>0)
    {
        if(Re<1e-5)
            error_msg("Reynolds number %g is too small.\n", (double)Re);
        switch(flags.advection)
        {
            case adv_explicit:
                PrintLog("Advection is treated explicitly (%u). Re = %g\n",
                    flags.advection, (double)Re);
                break;
            case adv_rk45:
                error_msg("Advection by Runge-Kutta 4-5 is not available\n");
                PrintLog("Advection is treated by rk45 (%u). Re = %g\n",
                    flags.advection, (double)Re);
                break;
            case adv_char:
                if(flags.moving_grid==0)
                {
                    warning_msg("The method of characteristics is only "
                        "available for moving grids.\n");
                    flags.advection=adv_semi;
                    /* ATTENTION: do not put `break' here and make sure adv_semi is next */
                }
                else
                {
                    PrintLog("Advection is treated by characteristics (%u). Re = %g\n",
                        flags.advection, (double)Re);
                    adv_nodes=calloc(grid->NN, sizeof(node_t));
                    MakeInverse(grid);
                    break;
                }
            case adv_semi:
                PrintLog("Advection is treated semi-implicitly (%u). Re = %g\n",
                    flags.advection, (double)Re);
                break;
        }
    }
    else
    {
        if(Re<1e-5)
            Re=1.0;
        PrintLog("Advection is OFF. Re = %g\n", Re);
    }

}

#if FALSE

There is a change here

static void set_order_flags(void)
{
    switch(flags.time_order)
    {
        case 0:
            PrintLog("Setting default time integration order 1.\n");
            flags.time_order=1;
            /* !!! do not put "break" here !!! */
        case 1:
            mass_coef=1.0;
            stif_coef=TStep;
            switch(flags.advection)
            {
                case adv_none: advm_coef=0.0;
                    break; /* no advection (Stokes) */
                case adv_explicit: advm_coef=0.0;
                    break; /* explicit advection */
                case adv_semi: advm_coef=TStep;
                    break; /* semi-implicit advection */
                case adv_char: advm_coef=0.0;
                    break; /* characteristics advection */
                case adv_rk45: advm_coef=0.0;
                    break; /* runge-kutta */
            }
            PrintLog("Time integration order 1. "
                "   mass_coef=%g    stif_coef=%g    advm_coef=%g.\n",
                (double)mass_coef, (double)stif_coef, (double)advm_coef);
            break;
        case 2:
        case 3:
            error_msg("Time integration order %d not implemented.\n",
                flags.time_order);
    }
}
#endif

